import pytest
import numpy as np
import openpnm as op
import numpy.testing as nt


class TransportTest:

    def setup_class(self):
        self.net = op.network.Cubic(shape=[9, 9, 9])
        self.net.add_model_collection(
            op.models.collections.geometry.spheres_and_cylinders)
        self.net.regenerate_models()
        self.phase = op.phase.Air(network=self.net)
        self.phase.add_model_collection(
            op.models.collections.physics.basic)
        self.phase.regenerate_models()
        self.phase['pore.mole_fraction'] = 0
        self.phase['throat.diffusive_conductance'] = 1.0

    def test_undefined_elements(self):
        net = op.network.Cubic(shape=[3, 3, 3])
        phase = op.phase.Phase(network=net)
        phase["throat.conductance"] = 1.0
        alg = op.algorithms.Transport(network=net, phase=phase)
        alg.settings._update({"quantity": "pore.concentration",
                              "conductance": "throat.conductance"})

    def test_remove_boundary_conditions(self):
        alg = op.algorithms.Transport(network=self.net, phase=self.phase)
        alg.set_value_BC(pores=self.net.pores('top'), values=1)
        alg.set_value_BC(pores=self.net.pores('bottom'), values=0)
        assert np.sum(np.isfinite(alg['pore.bc.value'])) > 0
        alg.set_value_BC(pores=self.net.pores('top'), mode='remove')
        assert np.sum(np.isfinite(alg['pore.bc.value'])) > 0
        alg.set_value_BC(pores=self.net.pores('bottom'), mode='remove')
        assert np.sum(np.isfinite(alg['pore.bc.value'])) == 0

    def test_generic_transport(self):
        alg = op.algorithms.Transport(network=self.net, phase=self.phase)
        alg.settings['conductance'] = 'throat.diffusive_conductance'
        alg.settings['quantity'] = 'pore.mole_fraction'
        alg.set_value_BC(pores=self.net.pores('top'), values=1)
        alg.set_value_BC(pores=self.net.pores('bottom'), values=0)
        alg.run()
        from openpnm.algorithms._solution import SteadyStateSolution
        quantity = alg.settings['quantity']
        # Ensure solution object is attached to the algorithm
        assert isinstance(alg.soln[quantity], SteadyStateSolution)

    def test_ill_defined_topology(self):
        net = op.network.Cubic(shape=[5, 1, 1])
        op.topotools.trim(net, pores=[2])
        phase = op.phase.Phase(network=net)
        phase["throat.diffusive_conductance"] = 1.0
        alg = op.algorithms.Transport(network=net, phase=phase)
        alg.settings['conductance'] = 'throat.diffusive_conductance'
        alg.settings['quantity'] = 'pore.mole_fraction'
        alg.set_value_BC(pores=net.pores('left'), values=1)
        with pytest.raises(Exception):
            alg.run()

    def test_linear_system_with_nans_or_infs(self):
        net = op.network.Cubic(shape=[5, 1, 1])
        phase = op.phase.Phase(network=net)
        phase["throat.diffusive_conductance"] = 1.0
        alg = op.algorithms.Transport(network=net, phase=phase)
        alg.settings['conductance'] = 'throat.diffusive_conductance'
        alg.settings['quantity'] = 'pore.mole_fraction'
        alg.set_value_BC(pores=net.pores('left'), values=1)
        # A contains nan
        with pytest.raises(Exception):
            phase["throat.diffusive_conductance"][1] = np.inf
            alg.run()
        phase["throat.diffusive_conductance"] = 1.0
        # A contains inf
        with pytest.raises(Exception):
            phase["throat.diffusive_conductance"][1] = np.inf
            alg.run()
        phase["throat.diffusive_conductance"] = 1.0
        # b contains nan
        with pytest.raises(Exception):
            phase["throat.diffusive_conductance"][1] = np.inf
            alg.run()
        phase["throat.diffusive_conductance"] = 1.0
        # b contains inf
        with pytest.raises(Exception):
            phase["throat.diffusive_conductance"][1] = np.inf
            alg.run()

    def test_ill_defined_settings(self):
        net = op.network.Cubic(shape=[5, 1, 1])
        phase = op.phase.Phase(network=net)
        phase["throat.diffusive_conductance"] = 1.0
        alg = op.algorithms.Transport(network=net, phase=phase)
        alg.settings['conductance'] = 'throat.diffusive_conductance'
        alg.settings['quantity'] = 'pore.mole_fraction'
        alg.set_value_BC(pores=net.pores('left'), values=1)
        # conductance is not defined
        with pytest.raises(Exception):
            alg.settings['conductance'] = None
            alg.run()
        alg.settings['conductance'] = 'throat.diffusive_conductance'
        # quantity is not defined
        with pytest.raises(Exception):
            alg.settings['quantity'] = None
            alg.run()
        alg.settings['quantity'] = 'pore.mole_fraction'
        # phase is not defined
        with pytest.raises(Exception):
            alg.settings['phase'] = None
            alg.run()

    def test_two_value_conditions(self):
        alg = op.algorithms.Transport(network=self.net,
                                             phase=self.phase)
        alg.settings['conductance'] = 'throat.diffusive_conductance'
        alg.settings['quantity'] = 'pore.mole_fraction'
        alg.set_value_BC(pores=self.net.pores('top'), values=1)
        alg.set_value_BC(pores=self.net.pores('bottom'), values=0)
        alg.run()
        x = [0.0, 0.125, 0.25, 0.375, 0.5, 0.625, 0.75, 0.875, 1.0]
        y = np.unique(np.around(alg['pore.mole_fraction'], decimals=3))
        assert np.all(x == y)

    def test_one_value_one_rate(self):
        alg = op.algorithms.Transport(network=self.net, phase=self.phase)
        alg.settings['conductance'] = 'throat.diffusive_conductance'
        alg.settings['quantity'] = 'pore.mole_fraction'
        alg.set_rate_BC(pores=self.net.pores('bottom'), rates=1)
        alg.set_value_BC(pores=self.net.pores('top'), values=0)
        alg.run()
        x = [0., 1., 2., 3., 4., 5., 6., 7., 8.]
        y = np.unique(np.around(alg['pore.mole_fraction'], decimals=3))
        assert np.all(x == y)

    def test_set_value_bc_where_rate_is_already_set_mode_merge(self):
        alg = op.algorithms.Transport(network=self.net, phase=self.phase)
        alg.settings['conductance'] = 'throat.diffusive_conductance'
        alg.settings['quantity'] = 'pore.mole_fraction'
        alg.set_rate_BC(pores=[0, 1], rates=1, mode='add')
        with pytest.raises(Exception):
            alg.set_value_BC(pores=[1, 2], values=0, mode='add')
        assert np.isfinite(alg['pore.bc.rate']).sum() == 2
        assert np.isfinite(alg['pore.bc.value']).sum() == 0
        alg.set_rate_BC(pores=[0, 1], rates=1, mode='overwrite')
        alg.set_value_BC(pores=[1, 2], values=0, mode='overwrite')
        assert np.isfinite(alg['pore.bc.rate']).sum() == 1
        assert np.isfinite(alg['pore.bc.value']).sum() == 2

    def test_cache(self):
        alg = op.algorithms.Transport(network=self.net, phase=self.phase)
        alg.settings['conductance'] = 'throat.diffusive_conductance'
        alg.settings['quantity'] = 'pore.mole_fraction'
        alg.set_rate_BC(pores=self.net.pores('bottom'), rates=1)
        alg.set_value_BC(pores=self.net.pores('top'), values=0)
        alg.settings["cache"] = True
        alg.run()
        x_before = alg["pore.mole_fraction"].mean()
        self.phase["throat.diffusive_conductance"][1] = 50.0
        alg.run()
        x_after = alg["pore.mole_fraction"].mean()
        # When cache is True, A is not recomputed, hence x == y
        assert x_before == x_after
        alg.settings["cache"] = False
        alg.run()
        x_after = alg["pore.mole_fraction"].mean()
        # When cache is False, A must be recomputed, hence x!= y
        assert x_before != x_after
        # Revert back changes to objects
        self.setup_class()

    def test_rate_single_pore(self):
        alg = op.algorithms.ReactiveTransport(network=self.net,
                                              phase=self.phase)
        alg.settings['conductance'] = 'throat.diffusive_conductance'
        alg.settings['quantity'] = 'pore.mole_fraction'
        pores = self.net.pores("left")
        alg.set_rate_BC(pores=pores, rates=1.235*np.ones(pores.size))
        alg.set_value_BC(pores=self.net.pores("right"), values=0.0)
        alg.run()
        rate = alg.rate(pores=self.net.pores("right"))[0]
        assert np.isclose(rate, -1.235*self.net.pores("right").size)
        # Net rate must always be zero at steady state conditions
        assert np.isclose(alg.rate(pores=self.net.Ps), 0.0)

    def test_rate_multiple_pores(self):
        alg = op.algorithms.Transport(network=self.net, phase=self.phase)
        alg.settings['conductance'] = 'throat.diffusive_conductance'
        alg.settings['quantity'] = 'pore.mole_fraction'
        alg.set_rate_BC(pores=[0, 1, 2, 3], rates=1.235)
        alg.set_rate_BC(pores=[5, 6, 19, 35, 0], rates=3.455, mode='overwrite')
        # Pore 0 is assigned two rate BCs, only the most recent will be kept
        alg.set_value_BC(pores=[50, 51, 52, 53], values=0.0)
        alg.run()
        rate = alg.rate(pores=[50, 51, 52, 53])[0]
        # 3 and 5 are number of pores in each rate BC
        assert np.isclose(rate, -(1.235*3 + 3.455*5))
        # Net rate must always be zero at steady state conditions
        assert np.isclose(alg.rate(pores=self.net.Ps), 0.0)

    def test_rate_multiple_values(self):
        alg = op.algorithms.Transport(network=self.net,
                                             phase=self.phase)
        alg.settings['conductance'] = 'throat.diffusive_conductance'
        alg.settings['quantity'] = 'pore.mole_fraction'
        alg.set_rate_BC(pores=[0, 1, 2, 3], rates=[0, 3.5, 0.4, -12])
        alg.set_value_BC(pores=[50, 51, 52, 53], values=0.0)
        alg.run()
        rate_individual = alg.rate(pores=[0, 1, 2, 3], mode='single')
        rate_net = alg.rate(pores=[0, 1, 2, 3], mode='group')[0]
        nt.assert_allclose(rate_individual, [0, 3.5, 0.4, -12], atol=1e-10)
        nt.assert_allclose(rate_net, sum([0, 3.5, 0.4, -12]))

    # def test_rate_Nt_by_2_conductance(self):
    #     net = op.network.Cubic(shape=[1, 6, 1])
    #     net.add_model_collection(
    #         op.models.collections.geometry.spheres_and_cylinders)
    #     net.regenerate_models()
    #     air = op.phase.Air(network=net)
    #     water = op.phase.Water(network=net)
    #     m = op.contrib.MultiPhase(network=net, phases=[air, water])
    #     m.set_occupancy(phase=air, pores=[0, 1, 2])
    #     m.set_occupancy(phase=water, pores=[3, 4, 5])
    #     const = op.models.misc.constant
    #     K_water_air = 0.5
    #     m.set_binary_partition_coef(phases=[water, air], model=const, value=K_water_air)
    #     m._set_automatic_throat_occupancy()
    #     alg = op.algorithms.Transport(network=net, phase=m)
    #     alg.settings['conductance'] = 'throat.diffusive_conductance'
    #     alg.settings['quantity'] = 'pore.mole_fraction'
    #     alg.set_rate_BC(pores=0, values=1.235)
    #     alg.set_value_BC(pores=5, values=0.0)
    #     alg.run()
    #     rate = alg.rate(pores=5)[0]
    #     assert np.isclose(rate, -1.235)
    #     # Rate at air-water interface throat (#2) must match imposed rate
    #     rate = alg.rate(throats=2)[0]
    #     assert np.isclose(rate, 1.235)
    #     # Rate at interface pores (#2 @ air-side, #3 @ water-side) must be 0
    #     rate_air_side = alg.rate(pores=2)[0]
    #     rate_water_side = alg.rate(pores=3)[0]
    #     assert np.isclose(rate_air_side, 0.0)
    #     assert np.isclose(rate_water_side, 0.0)
    #     # Net rate must always be zero at steady state conditions
    #     assert np.isclose(alg.rate(pores=net.Ps), 0.0)

    # def test_reset_settings_and_data(self):
    #     alg = op.algorithms.Transport(network=self.net,
    #                                          phase=self.phase)
    #     alg.settings['conductance'] = 'throat.diffusive_conductance'
    #     alg.settings['quantity'] = 'pore.mole_fraction'
    #     alg.set_rate_BC(pores=self.net.pores('bottom'), values=1)
    #     alg.set_value_BC(pores=self.net.pores('top'), values=0)
    #     alg.run()
    #     assert ~np.all(np.isnan(alg['pore.bc.value']))
    #     assert ~np.all(np.isnan(alg['pore.bc.rate']))
    #     assert 'pore.mole_fraction' in alg.keys()
    #     alg.reset(bcs=True, results=False)
    #     assert np.all(np.isnan(alg['pore.bc.value']))
    #     assert np.all(np.isnan(alg['pore.bc.rate']))
    #     assert 'pore.mole_fraction' in alg.keys()
    #     alg.reset(bcs=True, results=True)
    #     assert 'pore.mole_fraction' not in alg.keys()
    #     alg.set_rate_BC(pores=self.net.pores('bottom'), values=1)
    #     alg.set_value_BC(pores=self.net.pores('top'), values=0)
    #     alg.run()

    # def test_reset_actual_results(self):
    #     alg = op.algorithms.Transport(network=self.net,
    #                                          phase=self.phase)
    #     alg.settings['conductance'] = 'throat.diffusive_conductance_temp'
    #     self.phase['throat.diffusive_conductance_temp'] = 1.0
    #     alg.settings['quantity'] = 'pore.concentration'
    #     alg.set_value_BC(pores=self.net.pores('bottom'), values=1)
    #     alg.set_value_BC(pores=self.net.pores('top'), values=0)
    #     alg.run()
    #     m1 = alg.rate(pores=self.net.pores('top'))
    #     m2 = -alg.rate(pores=self.net.pores('bottom'))
    #     # This should pass because the alg has only run once
    #     np.testing.assert_allclose(m1, m2)
    #     # Now adjust conductance values and re-run
    #     self.phase['throat.diffusive_conductance_temp'][[0, 1, 2]] *= 0.1
    #     alg.run()
    #     m1 = alg.rate(pores=self.net.pores('top'))
    #     m2 = -alg.rate(pores=self.net.pores('bottom'))
    #     # The mass won't balance, so the same test will fail
    #     with pytest.raises(AssertionError):
    #         np.testing.assert_allclose(m1, m2)
    #     # Now use reset method
    #     alg.reset()
    #     alg.run()
    #     m1 = alg.rate(pores=self.net.pores('top'))
    #     m2 = -alg.rate(pores=self.net.pores('bottom'))
    #     # Now this will pass again
    #     np.testing.assert_allclose(m1, m2)

    # def test_validate_data_health(self):
    #     alg = op.algorithms.Transport(network=self.net,
    #                                          phase=self.phase)
    #     alg.settings['conductance'] = 'throat.diffusive_conductance'
    #     alg.settings['quantity'] = 'pore.concentration'
    #     alg.settings['cache'] = False
    #     alg.set_value_BC(pores=self.net.pores('top'), values=1)
    #     alg.set_value_BC(pores=self.net.pores('bottom'), values=0)
    #     # Check if the method can catch NaNs in data
    #     self.phase['throat.diffusive_conductance'][0] = np.nan
    #     with pytest.raises(Exception):
    #         alg.run()
    #     mod = op.models.misc.from_neighbor_pores
    #     self.phase["pore.seed"] = np.nan
    #     self.phys.add_model(propname="throat.diffusive_conductance", model=mod,
    #                         prop="pore.seed", ignore_nans=False)
    #     with pytest.raises(Exception):
    #         alg.run()
    #     self.phase["pore.seed"] = 1
    #     self.phase.regenerate_models(propnames="throat.diffusive_conductance")
    #     # Check if the method can catch unhealthy topology
    #     Ts = self.net.find_neighbor_throats(pores=0)
    #     op.topotools.trim(self.net, throats=Ts)
    #     with pytest.raises(Exception):
    #         alg.run()
    #     # Reset network back to original
    #     self.setup_class()

    # def test_total_rate(self):
    #     alg = op.algorithms.GenericTransport(network=self.net,
    #                                          phase=self.phase)
    #     h = op.utils.check_network_health(self.net)
    #     op.topotools.trim(self.net, pores=h['disconnected_pores'])
    #     alg.settings['conductance'] = 'throat.diffusive_conductance'
    #     alg.settings['quantity'] = 'pore.mole_fraction'
    #     alg.set_rate_BC(pores=[0, 1, 2, 3], total_rate=1)
    #     alg.set_value_BC(pores=[50, 51, 52, 53], values=0.0)
    #     alg.run()
    #     rate_individual = alg.rate(pores=[0, 1, 2, 3], mode='single')
    #     nt.assert_allclose(rate_individual, [0.25, 0.25, 0.25, 0.25],
    #                        atol=1e-10)
    #     # test exceptions that come from adding total_rate feature
    #     with pytest.raises(Exception):
    #         alg.set_rate_BC(pores=[0, 1, 2, 3],
    #                         total_rate=[0.25, 0.25, 0.25, 0.25])
    #     with pytest.raises(Exception):
    #         alg.set_rate_BC(pores=[0, 1, 2, 3], rates=1, total_rate=1)

    # def test_network_continuity(self):
    #     net = op.network.Cubic([5, 1, 1])
    #     op.topotools.trim(network=net, pores=[2])
    #     phase = op.phase.Phase(network=net)
    #     phase['throat.diffusive_conductance'] = 1.0
    #     alg = op.algorithms.FickianDiffusion(network=net, phase=phase)
    #     alg.set_value_BC(pores=0, values=1)
    #     with pytest.raises(Exception):
    #         alg.run()
    #     alg.set_value_BC(pores=3, values=0)
    #     alg.run()

    def test_x0_is_nan(self):
        alg = op.algorithms.Transport(network=self.net,
                                      phase=self.phase)
        alg.settings['conductance'] = 'throat.diffusive_conductance'
        alg.settings['quantity'] = 'pore.mole_fraction'
        alg.set_value_BC(pores=self.net.pores('top'), values=1)
        alg.set_value_BC(pores=self.net.pores('bottom'), values=0)
        x0 = np.zeros(alg.Np)
        x0[5] = np.nan
        with pytest.raises(Exception):
            alg.run(x0=x0)

    def test_get_source_list(self):
        alg = op.algorithms.Transport(network=self.net,
                                      phase=self.phase)
        assert alg['pore.source'] == {}
        with pytest.raises(KeyError):
            alg['pore.source_blah'] == {}

    def teardown_class(self):
        ws = op.Workspace()
        ws.clear()


if __name__ == '__main__':

    t = TransportTest()
    t.setup_class()
    self = t
    for item in t.__dir__():
        if item.startswith('test'):
            print(f'Running test: {item}')
            t.__getattribute__(item)()
